#!/bin/bash

# setup_root_ssh_key.sh
# 多系统支持：为root用户配置SSH密钥登录，支持交互式用户复制

# 设置颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

info() { echo -e "${BLUE}[INFO]${NC} $1"; }
success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; }

# 检查是否以root运行
if [ "$EUID" -ne 0 ]; then
    error "此脚本必须使用 root 权限运行！"
    echo "请使用: sudo bash $0"
    exit 1
fi

# 检测系统类型
detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS_NAME=$NAME
        OS_ID=$ID
        OS_VERSION=$VERSION_ID
    elif [ -f /etc/redhat-release ]; then
        OS_NAME=$(cat /etc/redhat-release)
        OS_ID="rhel"
    else
        OS_NAME="Unknown"
        OS_ID="unknown"
    fi
    info "检测到系统: $OS_NAME ($OS_ID)"
}

# 安装必要工具
install_ssh_tools() {
    case $OS_ID in
        ubuntu|debian)
            if ! command -v ssh-keygen &> /dev/null; then
                info "安装 openssh-client..."
                apt update && apt install -y openssh-client
            fi
            ;;
        centos|rhel|fedora|almalinux|rocky)
            if ! command -v ssh-keygen &> /dev/null; then
                info "安装 openssh-clients..."
                yum install -y openssh-clients 2>/dev/null || dnf install -y openssh-clients
            fi
            ;;
        *)
            warning "未知系统类型，跳过工具安装"
            ;;
    esac
}

# 配置防火墙
configure_firewall() {
    case $OS_ID in
        ubuntu|debian)
            if command -v ufw &> /dev/null && ufw status | grep -q "Status: active"; then
                if ! ufw status | grep -q "22.*ALLOW"; then
                    info "配置UFW防火墙允许SSH..."
                    ufw allow ssh
                fi
            fi
            ;;
        centos|rhel|fedora|almalinux|rocky)
            if systemctl is-active --quiet firewalld; then
                if ! firewall-cmd --list-all | grep -q "ssh"; then
                    info "配置Firewalld允许SSH..."
                    firewall-cmd --permanent --add-service=ssh
                    firewall-cmd --reload
                fi
            fi
            ;;
    esac
}

# 修复SELinux上下文
fix_selinux_context() {
    if command -v restorecon &> /dev/null && command -v sestatus &> /dev/null; then
        if sestatus | grep -q "SELinux status:.*enabled"; then
            info "修复SELinux上下文..."
            restorecon -R -v /root/.ssh 2>/dev/null
        fi
    fi
}

# 获取服务重载命令
get_reload_command() {
    case $OS_ID in
        ubuntu|debian|centos|rhel|fedora|almalinux|rocky)
            if systemctl is-active --quiet ssh &> /dev/null; then
                echo "systemctl reload ssh"
            elif systemctl is-active --quiet sshd &> /dev/null; then
                echo "systemctl reload sshd"
            else
                echo "service ssh reload"
            fi
            ;;
        *)
            echo "service ssh reload"
            ;;
    esac
}

# 获取日志文件路径
get_log_file() {
    case $OS_ID in
        ubuntu|debian)
            echo "/var/log/auth.log"
            ;;
        centos|rhel|fedora|almalinux|rocky)
            echo "/var/log/secure"
            ;;
        *)
            echo "/var/log/messages"
            ;;
    esac
}

# 询问是否复制到普通用户
ask_copy_to_user() {
    while true; do
        read -p "是否要复制密钥到普通用户目录？(y/N): " -n 1 -r
        echo
        case $REPLY in
            [Yy])
                return 0
                ;;
            [Nn]|"")
                return 1
                ;;
            *)
                echo "请输入 y 或 n"
                ;;
        esac
    done
}

# 获取普通用户名（带重试询问）
get_username() {
    local username
    while true; do
        read -p "请输入普通用户名: " username
        username=$(echo "$username" | tr -d ' ')  # 去除空格
        
        if [ -z "$username" ]; then
            warning "用户名不能为空"
            continue
        fi
        
        if id "$username" &>/dev/null; then
            echo "$username"
            return 0
        else
            warning "用户 $username 不存在！"
            
            while true; do
                read -p "是否重试？(Y/n): " -n 1 -r
                echo
                case $REPLY in
                    [Yy]|"")
                        break  # 继续外层循环重试
                        ;;
                    [Nn])
                        return 1  # 返回失败
                        ;;
                    *)
                        echo "请输入 y 或 n"
                        ;;
                esac
            done
        fi
    done
}

# 复制密钥到普通用户
copy_to_user() {
    local username=$1
    local user_ssh_dir="/home/$username/root_ssh_backup"
    
    info "正在为普通用户 $username 创建密钥备份..."
    mkdir -p "$user_ssh_dir"
    
    if cp "$KEY_PATH" "$user_ssh_dir/root_key" && cp "$KEY_PATH.pub" "$user_ssh_dir/root_key.pub"; then
        chown -R "$username:$username" "$user_ssh_dir"
        chmod 700 "$user_ssh_dir"
        chmod 600 "$user_ssh_dir/root_key"
        chmod 644 "$user_ssh_dir/root_key.pub"
        success "密钥已复制到普通用户目录: $user_ssh_dir"
        return 0
    else
        error "复制密钥到用户目录失败！"
        
        while true; do
            read -p "是否重试复制？(Y/n): " -n 1 -r
            echo
            case $REPLY in
                [Yy]|"")
                    return 2  # 需要重试
                    ;;
                [Nn])
                    return 1  # 不重试
                    ;;
                *)
                    echo "请输入 y 或 n"
                    ;;
            esac
        done
    fi
}

# 配置SSH服务
configure_ssh() {
    local sshd_config="/etc/ssh/sshd_config"
    
    # 备份原配置
    cp "$sshd_config" "$sshd_config.backup.$(date +%Y%m%d%H%M%S)"
    
    info "正在配置SSH服务..."
    
    # 修改配置
    sed -i 's/^#*PermitRootLogin.*/PermitRootLogin yes/' "$sshd_config"
    sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication yes/' "$sshd_config"
    sed -i 's/^#*PubkeyAuthentication.*/PubkeyAuthentication yes/' "$sshd_config"
    
    # 确保配置存在
    grep -q "^PermitRootLogin" "$sshd_config" || echo "PermitRootLogin yes" >> "$sshd_config"
    grep -q "^PasswordAuthentication" "$sshd_config" || echo "PasswordAuthentication yes" >> "$sshd_config"
    grep -q "^PubkeyAuthentication" "$sshd_config" || echo "PubkeyAuthentication yes" >> "$sshd_config"
    
    # 重载SSH服务
    local reload_cmd=$(get_reload_command)
    if eval "$reload_cmd"; then
        success "SSH配置已更新并重载成功！"
    else
        warning "SSH服务重载失败，尝试其他方式..."
        # 尝试其他重载方式
        if systemctl restart ssh 2>/dev/null || systemctl restart sshd 2>/dev/null || service ssh restart 2>/dev/null; then
            success "SSH服务重启成功！"
        else
            error "SSH服务重启失败，请手动检查！"
        fi
    fi
}

# 获取密码（带确认）
get_password() {
    local password1 password2
    while true; do
        read -s -p "请输入root密钥的密码（直接回车则为空）: " password1
        echo
        if [ -z "$password1" ]; then
            # 返回空字符串
            echo ""
            return 0
        fi
        
        read -s -p "请再次输入密码确认: " password2
        echo
        
        if [ "$password1" = "$password2" ]; then
            echo "$password1"
            return 0
        else
            error "两次输入的密码不一致，请重新输入！"
        fi
    done
}

# 显示总结信息
show_summary() {
    local username=$1
    local user_ssh_dir=$2
    local server_ip=$(hostname -I | awk '{print $1}')
    local log_file=$(get_log_file)
    
    echo
    echo "================================================================"
    success "SSH密钥设置完成！"
    echo
    info "系统信息:"
    echo "操作系统: $OS_NAME"
    echo "系统ID: $OS_ID"
    echo "服务器IP: ${server_ip}"
    echo "日志文件: ${log_file}"
    echo
    info "密钥信息:"
    echo "root私钥: ${KEY_PATH}"
    echo "root公钥: ${KEY_PATH}.pub"
    
    if [ -n "$username" ] && [ -d "/home/$username/root_ssh_backup" ]; then
        echo "用户备份: /home/$username/root_ssh_backup/"
        echo
        info "下载命令:"
        echo -e "${GREEN}scp ${username}@${server_ip}:/home/${username}/root_ssh_backup/root_key ./${NC}"
    fi
    
    echo
    info "连接命令:"
    echo -e "${GREEN}ssh -i ${KEY_PATH} root@${server_ip}${NC}"
    echo
    info "测试命令:"
    echo -e "${GREEN}ssh -i ${KEY_PATH} root@localhost 'echo SSH密钥登录测试成功'${NC}"
    echo
    info "当前SSH配置:"
    echo "✓ Root登录: 允许"
    echo "✓ 密钥认证: 允许" 
    echo "✓ 密码认证: 允许"
    echo "================================================================"

    # 安全建议
    echo
    warning "安全建议:"
    echo "1. 测试成功后建议禁用密码登录:"
    echo "   sed -i 's/^PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config"
    echo "   然后重启SSH服务"
    echo "2. 查看SSH日志: tail -f ${log_file} | grep ssh"
    echo "3. 设置强root密码: passwd root"
    if command -v sestatus &> /dev/null; then
        echo "4. SELinux状态: $(sestatus | grep 'SELinux status' | cut -d: -f2 | tr -d ' ')"
    fi
    if [ -n "$username" ] && [ -d "/home/$username/root_ssh_backup" ]; then
        echo "5. 已为用户 ${username} 创建密钥备份，请及时下载并妥善保管"
    fi
    echo "6. 注意: 旧的SSH密钥已失效，只有新生成的密钥可以登录"
}

# 主执行流程
main() {
    # 全局变量
    local KEY_TYPE="ed25519"
    local KEY_PATH="/root/.ssh/id_$KEY_TYPE"
    local USERNAME=""
    local USER_SSH_DIR=""
    local PASSPHRASE=""
    
    # 显示欢迎信息
    echo "================================================================"
    echo "            SSH Root密钥配置工具 - 多系统支持版"
    echo "================================================================"
    
    # 检测系统类型
    detect_os
    
    # 安装必要工具
    install_ssh_tools
    
    # 创建root的.ssh目录
    mkdir -p /root/.ssh
    chmod 700 /root/.ssh

    # 检查是否已存在密钥
    if [ -f "$KEY_PATH" ]; then
        warning "发现已存在的私钥文件: $KEY_PATH"
        read -p "是否覆盖？(y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            info "操作已取消。"
            exit 0
        fi
        rm -f "$KEY_PATH" "$KEY_PATH.pub"
    fi

    # 生成新的SSH密钥对
    info "正在生成新的 $KEY_TYPE 密钥对..."
    
    # 获取密码（带确认）
    PASSPHRASE=$(get_password)
    
    if [ -z "$PASSPHRASE" ]; then
        ssh-keygen -t "$KEY_TYPE" -f "$KEY_PATH" -q -N ""
        success "已生成无密码的SSH密钥对！"
    else
        ssh-keygen -t "$KEY_TYPE" -f "$KEY_PATH" -q -N "$PASSPHRASE"
        success "已生成带密码的SSH密钥对！"
    fi

    if [ $? -ne 0 ]; then
        error "密钥生成失败！"
        exit 1
    fi

    # 重要修改：清空并替换authorized_keys，使旧密钥立即失效
    info "正在配置公钥授权（旧的密钥将失效）..."
    cat "$KEY_PATH.pub" > /root/.ssh/authorized_keys
    chmod 600 /root/.ssh/authorized_keys
    success "公钥已成功配置，旧的SSH密钥已失效！"

    # 询问是否复制到普通用户
    if ask_copy_to_user; then
        # 直接调用用户名获取函数
        USERNAME=$(get_username)
        
        if [ -n "$USERNAME" ] && id "$USERNAME" &>/dev/null; then
            # 尝试复制密钥
            while true; do
                copy_to_user "$USERNAME"
                local copy_result=$?
                
                case $copy_result in
                    0)  # 复制成功
                        USER_SSH_DIR="/home/$USERNAME/root_ssh_backup"
                        break
                        ;;
                    1)  # 用户选择不重试
                        warning "用户选择不重试复制操作，跳过复制步骤"
                        break
                        ;;
                    2)  # 需要重试复制
                        info "准备重试复制操作..."
                        continue
                        ;;
                esac
            done
        else
            warning "无法获取有效用户名，跳过复制步骤"
        fi
    else
        info "用户选择不复制到普通用户目录"
    fi

    # 配置SSH服务
    configure_ssh

    # 配置防火墙
    configure_firewall

    # 修复SELinux上下文
    fix_selinux_context

    # 显示最终信息
    show_summary "$USERNAME" "$USER_SSH_DIR"
}

# 执行主函数
main "$@"